'''
https://leetcode.cn/problems/shortest-path-with-alternating-colors/?envType=study-plan-v2&envId=graph-theory
'''
import heapq
from typing import List


class Solution:
    def shortestAlternatingPaths(self, n: int, redEdges: List[List[int]], blueEdges: List[List[int]]) -> List[int]:
        # graph[u][0][v]： u -> v 有红边
        # graph[u][1][v]： u -> v 有红蓝边
        graph = [[[], []] for _ in range(n)]
        for u, v in redEdges:
            graph[u][0].append(v)
        for u, v in blueEdges:
            graph[u][1].append(v)

        distance = [[float('inf'), float('inf')] for _ in range(n)]
        distance[0][0] = 0
        distance[0][1] = 0
        pq = [(0, 0, 0), (0, 1, 0)]  # (cost, red/blue, node)
        # q = [(0, 0, 0), (0, 1, 0)]  # (cost, red/blue, node)
        while pq:
            cost, color, u = heapq.heappop(pq)
            # cost, color, u = q.pop()
            next_color = 1 - color
            for v in graph[u][color]:
                if cost + 1 < distance[v][color]:
                    distance[v][color] = cost + 1
                    heapq.heappush(pq, (distance[v][color], next_color, v))
                    # q.append((distance[v][color], next_color, v))
        res = [0] * n
        for i in range(n):
            res[i] = min(distance[i])
            if res[i] == float('inf'):
                res[i] = -1
        return res
n = 5
red_edges = [[0,1],[1,2],[2,3],[3,4]]
blue_edges = [[1,2],[2,3],[3,1]]
print(Solution().shortestAlternatingPaths(n, red_edges, blue_edges))